# Kestrel: An XMPP-based Many-Task Computing Scheduler
# Author: Lance Stout <lancestout@gmail.com>
#
# Copyright 2009-2010 Clemson University
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


"""
XMPP client module.

The XMPP module provides an interface for sending and receiving
both XMPP messages and presence notifications.
"""

import sleekxmpp
import sleekxmpp.componentxmpp
import yaml


class Base(object):
    """
    Collection of XMPP event handlers common to both the Client and Component
    implementations.
    """

    def __init__(self):
        """
        Instantiate a group of basic XMPP connection handlers.
        """
        self.xmpp = None
        self.kernel = None
        self.has_roster = True
        self.jid = ''

    def disconnect(self, data=None):
        """Disconnect from the XMPP server."""
        self.xmpp.disconnect()
        self.kernel.event('system/offline')

    def on_start(self, event):
        """Fire system/available event once a connection has been made."""
        if self.has_roster:
            self.xmpp.getRoster()
        self.xmpp.sendPresence()
        self.kernel.event('system/online')
        self.kernel.event('system/available')

    def recv_status(self, presence):
        """Fire a presence event when a presence notification is received."""
        jid = str(presence['from'])
        bare = presence['from'].bare

        presence_map = {'chat': 'presence/available',
                        'available': 'presence/available',
                        'away': 'presence/busy',
                        'dnd': 'presence/busy',
                        'unavailable': 'presence/offline'}
        event_type = presence_map.get(presence['type'], 'unavailable')

        self.kernel.event(event_type, {'full': jid, 'bare': bare})

    def recv_message(self, msg):
        """
        Fire a message event when a message has been received.

        Currently, it is assumed that message bodies will be
        written in either YAML or JSON.
        """
        sender = str(msg['from'])
        packet = yaml.load(msg['body'])

        if msg['type'] == 'error':
            self.kernel.event('message/error', {'from': sender,
                                                'body': packet})
        else:
            self.kernel.event('message', {'from': sender, 'body': packet})

    def send_status(self, status):
        """
        Send presence notification that the system is either
        available or busy.

        Status can either be a pure broadcast, or directed
        to particular entity. For a directed presence, use
        the format {'to': 'user@example.com', 'type':'available'}.
        """

        to = None
        if type(status) is dict:
            to = status['to']
            status = status['type']

        presence_map = {'available': {'pshow': 'chat',
                                      'pstatus': 'Available'},
                        'busy': {'pshow': 'dnd',
                                 'pstatus': 'Running Task'},
                        'subscribe': {'ptype': 'subscribe'},
                        'unsubscribe': {'ptype': 'unsubscribe'},
                        'subscribed': {'ptype': 'subscribed'},
                        'unsubscribed': {'ptype': 'unsubscribed'}}

        presence = presence_map.get(status, None)
        if None not in [presence, to]:
            presence['pto'] = to

        state = self.kernel.config('system/status', 'available')
        state = presence_map[state]['pshow']

        if status == 'subscribe' \
                and to in self.xmpp.roster.keys() \
                and self.xmpp.roster[to]['subscription'] == 'both':
            self.kernel.event('system/access/approved', to)
            self.xmpp.sendPresence(pto=to, pshow=state)
            return

        if None not in [self.xmpp, presence]:
            if status in ['available', 'busy']:
                for jid in self.xmpp.roster.keys():
                    presence['pto'] = jid
                    self.xmpp.sendPresence(**presence)
            elif status in ['subscribe', 'unsubscribe',
                            'unsubscribe', 'unsubscribed']:
                self.xmpp.sendPresence(**presence)
                self.xmpp.sendPresence(pto=to, pshow=state)
            else:
                self.xmpp.sendPresence(**presence)

    def send_message(self, packet):
        """
        Send an XMPP message. There are no assumptions made about
        the content of the message.
        """
        if self.xmpp is not None:
            self.xmpp.sendMessage(mto=packet['to'],
                                  mbody=packet['body'],
                                  mfrom=self.jid)

    def change_roster(self, presence):
        """Fire a subscribe or unsubscribe event when a request is received."""
        jid = presence['from'].bare

        presence_map = {'subscribe': 'presence/subscribe',
                        'unsubscribe': 'presence/unsubsribe',
                        'subscribed': 'presence/subscribed',
                        'unsubscribed': 'presence/unsubscribed'}
        event_type = presence_map.get(presence['type'])

        self.kernel.event(event_type, jid)


class Client(Base):
    """
    A client XMPP module. Using this module should typically
    be done only for testing or small applications due to XMPP
    client scalability issues. For applications requiring
    thousands of entries in their roster, use an XMPP server
    component instead.

    Configuration entries for this module:
      XMPP/client/username
      XMPP/client/password
      XMPP/client/server
      XMPP/type

    Events listened for by this module:
      system/start
      system/shutdown
      message/send
      presence/send

    Events raised by this module:
      system/available
      system/online
      system/offline
      message
      presence/available
      presence/busy
      presence/offline
      presence/subscribe
      presence/subscribed
      presence/unsubscribe
      presence/unsubscribed
    """

    def __init__(self, kernel):
        """
        Instantiate an XMPP client module that will connect
        to the given server using the username and password
        as credentials.
        """
        self.xmpp = None
        self.kernel = kernel

        username = kernel.config('XMPP/client/username')
        self.server = kernel.config('XMPP/client/server')
        self.password = kernel.config('XMPP/client/password')
        self.jid = username + '@' + self.server
        self.has_roster = True

        self.events = {'system/start': self.connect,
                       'system/shutdown': self.disconnect,
                       'message/send': self.send_message,
                       'presence/send': self.send_status}

        self.kernel.register(self)

    def connect(self, data=None):
        """
        Start the XMPP module and establish a connection with the server.
        """
        self.kernel.set_config('XMPP/type', 'client')

        self.xmpp = sleekxmpp.ClientXMPP(self.jid, self.password)

        # Since the SleekXMPP library also uses an event engine, we
        # must provide event handlers to it which will in turn raise
        # events for our kernel to process.
        self.xmpp.add_event_handler("session_start", self.on_start)
        self.xmpp.add_event_handler("message", self.recv_message,
                                    threaded=True)
        self.xmpp.add_event_handler("got_online", self.recv_status)
        self.xmpp.add_event_handler("got_offline", self.recv_status)
        self.xmpp.add_event_handler("changed_status", self.recv_status)
        self.xmpp.add_event_handler("changed_subscription", self.change_roster)
        # SleekXMPP by default authorizes and subscribes to any subscription
        # request. Overriding these settings is needed in order to make
        # subscription behaviour configurable through application policy.
        self.xmpp.auto_authorize = None
        self.xmpp.auto_subscribe = False

        if self.xmpp.connect((self.server, 5222)):
            self.xmpp.process(threaded=False)
        else:
            self.kernel.log("ERROR", "Could not establish XMPP connection.")


class Component(Base):
    """
    A component XMPP module. Using this module should typically
    be done in applications with a large number of subscribed agents.

    Configuration entries for this module:
      XMPP/component/username
      XMPP/component/secret
      XMPP/component/server
      XMPP/component/port
      XMPP/type

    Events listened for by this module:
      system/start
      system/shutdown
      message/send
      presence/send

    Events raised by this module:
      system/available
      system/online
      system/offline
      message
      presence/available
      presence/busy
      presence/offline
      presence/subscribe
      presence/subscribed
      presence/unsubscribe
      presence/unsubscribed
    """

    def __init__(self, kernel):
        """
        Instantiate an XMPP client module that will connect
        to the given server using the username and password
        as credentials.
        """
        self.xmpp = None
        self.kernel = kernel

        username = kernel.config('XMPP/component/username')
        self.server = kernel.config('XMPP/component/server')
        self.secret = kernel.config('XMPP/component/secret')
        self.port = kernel.config('XMPP/component/port')
        self.jid = username + '.' + self.server
        self.has_roster = False

        self.events = {'system/start': self.connect,
                       'system/shutdown': self.disconnect,
                       'message/send': self.send_message,
                       'presence/send': self.send_status,
                       'system/roster/loaded': self.connect}

        self.kernel.register(self)

    def connect(self, roster=None):
        """
        Start the XMPP module and establish a connection with the server.
        """
        self.kernel.set_config('XMPP/type', 'component')

        if roster is None:
            roster = self.kernel.config('XMPP/roster', 'roster.tmp')
            self.kernel.event('system/roster/load', roster)
            return

        self.xmpp = sleekxmpp.componentxmpp.ComponentXMPP(self.jid,
                                                          self.secret,
                                                          self.server,
                                                          self.port)
        if roster is not None:
            self.xmpp.roster = roster

        # Since the SleekXMPP library also uses an event engine, we
        # must provide event handlers to it which will in turn raise
        # events for our kernel to process.
        self.xmpp.add_event_handler("session_start", self.on_start)
        self.xmpp.add_event_handler("message", self.recv_message,
                                    threaded=True)
        self.xmpp.add_event_handler("got_online", self.recv_status,
                                    threaded=True)
        self.xmpp.add_event_handler("got_offline", self.recv_status,
                                    threaded=True)
        self.xmpp.add_event_handler("changed_status", self.recv_status,
                                    threaded=True)
        self.xmpp.add_event_handler("changed_subscription", self.change_roster,
                                    threaded=True)
        # SleekXMPP by default authorizes and subscribes to any subscription
        # request. Overriding these settings is needed in order to make
        # subscription behaviour configurable through application policy.
        self.xmpp.auto_authorize = None
        self.xmpp.auto_subscribe = False

        if self.xmpp.connect():
            self.xmpp.process(threaded=False)
        else:
            self.kernel.log("ERROR", "Could not establish XMPP connection.")
